from datetime import datetime, timedelta, timezone
import json
import os
import pandas as pd
from tqdm import tqdm


# from pytypes import typechecked

dir_data = "./data"
file_market_map_codes = os.path.join(dir_data, "market_codes.json")
dir_df = os.path.join(dir_data, "dataframes")
file_code_market = os.path.join(dir_data, "code_market_info.csv")
file_all_stock_list = os.path.join(dir_data, "all_stock.txt")
file_sector_list = os.path.join(dir_data, "sector_list.json")
file_sector_to_code = os.path.join(dir_data, "sector_codes.json")

os.makedirs(dir_data, exist_ok=True)
os.makedirs(dir_df, exist_ok=True)


def get_code_market_info() -> pd.DataFrame:
    """
    :return: {'ExchangeID': 'DCE',
         'InstrumentID': 'eb00',
         'InstrumentName': '苯乙烯连续(2407)',
         'ProductID': 'eb',
         'ProductName': '',
         'ExchangeCode': None,
         'UniCode': None,
         'CreateDate': '0',
         'OpenDate': '20230727',
         'ExpireDate': 20240726,
         'PreClose': 9337.0,
         'SettlementPrice': 9435.0,
         'UpStopPrice': 10095.0,
         'DownStopPrice': 8775.0,
         'FloatVolume': 0.0,
         'TotalVolume': 0.0,
         'LongMarginRatio': 0.08,
         'ShortMarginRatio': 0.08,
         'PriceTick': 1.0,
         'VolumeMultiple': 5,
         'MainContract': 1,
         'LastVolume': 248082,
         'InstrumentStatus': 2147483647,
         'IsTrading': False,
         'IsRecent': False,
         'ProductTradeQuota': 0,
         'ContractTradeQuota': 0,
         'ProductOpenInterestQuota': 0,
         'ContractOpenInterestQuota': 24809}
    """
    if not os.path.exists(file_code_market):
        code_df = pd.read_csv(file_all_stock_list)
        code_market_list = code_df.iloc[:, 1].drop_duplicates().tolist()
        results = []
        from xtquant.xtdata import get_instrument_detail

        for code_market in tqdm(code_market_list, desc="get_instrument"):
            tmp = get_instrument_detail(code_market)

            if tmp:
                tmp["code_market"] = code_market
                tmp["market"] = code_market.split(".", 1)[-1]
                results.append(tmp)

        df = pd.DataFrame(results)
        df.to_csv(file_code_market)
        return df
    df = pd.read_csv(file_code_market, index_col=0)
    return df


def get_code_market_data(
    code_market: str,
    start: str = "",
    end: str = "",
    period: str = "1m",
    dividend_type="front_ratio",
    fill_data=True,
):
    dir_target = os.path.join(dir_df, code_market, period, dividend_type)
    os.makedirs(dir_target, exist_ok=True)
    file_target_df = os.path.join(dir_target, "kline.csv")
    file_time_range = os.path.join(dir_target, "info.json")
    if os.path.exists(file_time_range):
        with open(file_time_range) as f:
            info = json.load(f)
            exist_start = info["start"]
            exist_end = info["end"]
        if (not start or exist_start <= start) and (not end or exist_end >= end):
            df = pd.read_csv(file_target_df, index_col=0)
            df.index = df.index.astype(str)
            if start:
                start_time = datetime.strptime(start, "%Y%m%d").timestamp() * 1000
                df = df[df["time"] >= start_time]
            if end:
                end_time = (datetime.strptime(end, "%Y%m%d").timestamp()) * 1000
                df = df[df["time"] < end_time]
            return df

    from xtquant.xtdata import get_local_data

    dfs = get_local_data(
        field_list=[],
        stock_list=[code_market],
        period=period,
        start_time=start,
        end_time=end,
        count=-1,
        dividend_type=dividend_type,
        fill_data=fill_data,
        data_dir=None,
    )
    df = dfs[code_market]
    if not start:
        start = df.index.min()[:8]
    if not end:
        end = (
            datetime.strptime(df.index.max()[:8], "%Y%m%d") + timedelta(days=1)
        ).strftime("%Y%m%d")

    with open(file_time_range, "w") as f:
        info = {"start": start, "end": end}
        json.dump(info, f, ensure_ascii=False)
    df.to_csv(file_target_df)
    return df


def get_all_sector_list():
    if os.path.exists(file_sector_list):
        with open(file_sector_list) as f:
            sector_list = json.load(f)
        return sector_list

    from xtquant.xtdata import get_sector_list

    sectors = get_sector_list()
    with open(file_sector_list, "w") as f:
        json.dump(sectors, f, ensure_ascii=False)
    return sectors


def get_all_stock_list_in_sector(sector: str):
    if os.path.exists(file_sector_to_code):
        with open(file_sector_to_code) as f:
            sector_to_code = json.load(f)
            return sector_to_code[sector]

    from xtquant.xtdata import get_stock_list_in_sector

    sector_list = get_all_sector_list()
    results = {}
    for x in tqdm(sector_list, desc="获取sector结果"):
        results[x] = get_stock_list_in_sector(x)
    with open(file_sector_to_code, "w") as f:
        json.dump(results, f, ensure_ascii=False)
    return results[sector]


def get_baostock_list(xdate: str):
    """ "
    :param xdate: like 2020-12-31
    :return:
    """
    import baostock as bs

    result = bs.query_all_stock(xdate)
    df = result.get_data()
    return df


def get_baostock_data(
    stock_code, start_date, end_date, cache_dir: str = "./data", is_cover: bool = False
):
    import baostock as bs

    os.makedirs(cache_dir, exist_ok=True)

    df_file = os.path.join(cache_dir, f"{stock_code}.csv")
    if os.path.exists(df_file) and not is_cover:
        return pd.read_csv(df_file).query(
            f"'{start_date}' <= date and '{end_date}' >= date"
        )
    fields = (
        "date,code,open,high,low,close,volume,amount,"
        "adjustflag,turn,tradestatus,pctChg,peTTM,"
        "pbMRQ,psTTM,pcfNcfTTM,isST"
    )

    df = bs.query_history_k_data_plus(
        stock_code, fields, None, None, adjustflag="2"
    ).get_data()
    df.to_csv(df_file)
    return df.query(f"'{start_date}' <= date and '{end_date}' >= date")


def get_baostock_list(day: str = None):
    import baostock as bs
    import pandas as pd

    if not day:
        day = datetime.now().strftime("%Y-%m-%d")
    lg = bs.login()
    if lg.error_code != "0":
        print("login respond error_code:" + lg.error_code)
        print("login respond  error_msg:" + lg.error_msg)

    rs = bs.query_all_stock(day=day)
    if rs.error_code != "0":
        print("login respond error_code:" + rs.error_code)
        print("login respond  error_msg:" + rs.error_msg)

    data_list = []
    while (rs.error_code == "0") & rs.next():
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)

    bs.logout()
    return result


# @typechecked
def date_str8_to_sc_int(xt: str, tzinfo=timezone.utc) -> int:
    dt = datetime(
        year=int(xt[0:4]), month=int(xt[4:6]), day=int(xt[6:8]), tzinfo=tzinfo
    )
    return int(dt.timestamp())
